fix(executorch): support KV-cache aliased I/O in the TensorRT delegate - #4445

line 223:

Only inputs need this. Outputs are also bound positionally by the runtime,
but they are getitem(engine_node, idx) nodes whose index order equals the
engine output-binding order. ExecuTorch lowering can reorder delegate outputs
(arrange_graph_outputs moves buffer-mutation outputs ahead of user
outputs), but a TensorRT delegate partition is a functional inference engine
with no mutation outputs, so that pass is a no-op here and the output order is
preserved. If a TRT partition ever produced mutation outputs, outputs would
need the same node-identity reordering as inputs.

We need to account for the output order, or otherwise there is a mismatch

The engine appends the buffer mutation at the end of output, while the delegate prepends it

Runtime walkthrough

Assume:

TRT inputs:  [tokens, k_cache_in]
TRT outputs: [logits, k_cache_out]

ExecuTorch passes:

args[0] = tokens
args[1] = k_cache_in
args[2] = k_cache_mutation
args[3] = logits

After consuming the inputs, arg_idx == 2.

First output iteration

o    = 0
name = output_binding_names[0] = logits
arg  = args[2] = k_cache_mutation

The backend binds TensorRT's logits output to the cache-mutation output storage.

Second output iteration

o       = 1
name    = output_binding_names[1] = k_cache_out
out_arg = args[3] = logits

TensorRT correctly binds k_cache_out to k_cache_in for the in-place update, but the backend treats the logits EValue as its mutation output slot. Its reflect copy therefore writes the cache result into the logits output.

The result is effectively:

If shapes or capacities differ, execution may fail during resize/binding/enqueue. If they are compatible, it can run successfully while returning incorrect logits and corrupting the observable cache state.

The required fix is to reorder serialized output_binding_names into actual delegate-output order, analogous to _reorder_input_names_for_executorch.

I was not able to reproduce this. Seems like arrange_graph_outputs reorders the submodule outputs, output_specs, and the parent getitems, but it doesn't touch the call node's meta["val"]. node.args is permuted by fusion (so _reorder_input_names_for_executorch is still needed), but outputs go through meta["val"], which arrange_graph_outputs leaves alone.

If meta["val"] is being rearranged somewhere that would be a problem, but the fix would be different from how _reorder_input_names_for_executorch handles it.

Ok I spend a ton of time on this and it is quite surprising that the order was actually correct. But here is some findings:

_keep_mutated_buffers_above_delegate(exported_program)
https://github.com/Conarnar/TensorRT/blob/b9dbb306a4597eb4fb0c4ef0fdb546235d3924b1/py/torch_tensorrt/executorch/partitioner.py#L165
This function lifes mutated buffer to the executorch program level, therefore arrange_graph_outputs did not have any buffer_mutation, and therefore the order is perserved.

These tests verify that mutations get prepended in the ExportedProgram, which is one half of the contract. The other half — that the serialized output_binding_names line up with the delegate arg order the C++ backend indexes positionally — isn't covered here or anywhere else, because to_edge, arrange_graph_outputs, and preprocess are all outside the mocked boundary.

Could you add a test at that seam? Roughly:

1. Build a KV model with both a user output and an aliased cache output.
2. Run it through to_edge_transform_and_lower with the TRT partitioner.
3. Read the delegate's output-arg order from the lowered module's output specs.
4. Assert deserialize_engine(...).io_bindings output names match that order.

This needs no GPU and would fail today.

Aliased inputs can still take the host-staging path, which breaks caller-owned semantics.

If an aliased input's data pointer isn't CUDA-accessible, this falls through to the staging branch and binds engine->cached_input_ptrs[i] — a delegate-owned scratch buffer. The aliased output then binds to input_bind_ptrs[alias_in], i.e. that same scratch buffer, so the in-place KV update lands in delegate scratch rather than the caller's storage. On the next execute() the staging copy re-reads the caller's unchanged host buffer, so the update is silently lost. Recovery depends entirely on the reflect copy, which has its own problem (see the reflect comment).

Can you clarify what RFC 0003 §6.4 (and §7.4) is? I was not able to find any references for that.

Oh, here it is:
Require the caller's tensor itself to be device-resident. Aliasing only
works when the pointer bound to the input is the caller's real storage. Two
existing branches break that: the H2D fallback copies a host tensor into a
device staging buffer and binds that, and the zero-byte branch binds a
cached scratch allocation. Either would make the engine write the update
somewhere the caller never reads. For an alias-source input, gate on the
caller tensor being device/unified memory (is_cuda_accessible_ptr) with
non-zero size, and otherwise return Error::InvalidArgument ("aliased input
'%s' must be on GPU or unified memory") — i.e. reject before both the
staging and zero-byte paths, not only the H2D one.

Do we need to account for aliase I/O? Is there a race possible?

Yes, but only on the non-zero-copy reflect path: with a caller stream active and no end sync, a pending reflect into the delegate output could still be in flight when ExecuTorch's buffer-mutation copy_ reads it. Will handle it.

This doesn't cross-check the persisted alias map against the engine, and it accepts unknown kind values.

The Python runtime's _TRTEngine._reconcile_aliased_io treats getAliasedInputTensor as the source of truth for kv_cache_update aliases and preserves user ones as metadata-trusted. Here, a kind that is neither "kv_cache_update" nor "user" — a typo in the wire format, or a future kind written by a newer exporter — skips the shape check and gets registered as if it were a KV alias, which then binds two tensors to the same storage.

Could you mirror the Python behavior:

1. Reject unknown kinds with Error::InvalidProgram.
2. For kv_cache_update, compare the persisted ab.input against engine->getAliasedInputTensor(ab.output.c_str()) and error on disagreement.
3. Keep user as metadata-trusted after the shape check (TRT can't see those aliases).

Related: the parser leaves ab.kind empty when the "kind" key is absent, while the Python side defaults to "kv_cache_update". Once unknown kinds are rejected, that mismatch turns an old blob into a hard failure — worth defaulting to "kv_cache_update" in the parser to match.

This'd better be warnings or errors

Guessing you mean the two continues below.

Missing stream drain if a reflect copy fails, so the next call can touch a live context.

enqueueV3() has already succeeded by the time we get here, so the engine is running on the stream. If one of these cudaMemcpyAsync calls fails we return straight away, and engine->inflight_pending is only set further down (line 811), so it stays false.

Both later guards are gated on that flag:

• the next execute() (line 470) waits only if (engine->inflight_pending), and it does that before calling setInputShape / setTensorAddress
• the destructor (line 83) waits only if (inflight_pending) before destroying the context and freeing the staging buffers

So the flag says "nothing in flight" while TensorRT is still working. The next call is then free to reconfigure the context, which the comment at line 466 correctly says TensorRT forbids.

You already have exactly the right pattern a few lines below, at 802-808:

// Could not arm the completion marker; drain now so a later execute() or the
// destructor never reconfigures or frees exec_ctx while this enqueue runs.
(void)cudaStreamSynchronize(stream);
engine->inflight_pending = false;
return Error::InvalidProgram;

Same two lines would fix this branch.

Worth considering the more robust shape too: set the pending marker as soon as the first async work is submitted, rather than at the end of the happy path. Then any early return added later between enqueue and the end of the function inherits the protection instead of quietly reintroducing this.

Note the D2H branch at 786-793 looks like it has the same gap, but that one pre-dates this PR, so it is probably a separate fix.

Aliased models saved in the other two formats get no equivalent handling, and no warning.

This normalization runs only in the output_format == "executorch" branch. The exported_program branch (1116) and aot_inductor branch (1126) receive the same retrace=True program and save it as-is.

if output_format == "exported_program":     # no alias handling
elif output_format == "aot_inductor":       # no alias handling
elif output_format == "executorch":
    exp_program = _declare_aliased_kv_mutations_on_ep(exp_program)   # only here

I understand why it lives here: ExecuTorch has tag_constant_data, which freezes an undeclared buffer, and that pass is the reason this function exists. The other formats have no equivalent, so the harm is not the same.

But the outcome still differs per format in a way a user cannot see:

• exported_program: the engine still mutates the buffer in place at runtime, so the cache probably does update, but the saved signature does not declare the mutation. Anything downstream that trusts the signature gets a wrong answer.
• aot_inductor: this goes through inductor, and whether an in-place mutation on a custom op survives functionalization is not obvious. I could not find test coverage for that combination.

If aliased I/O is only intended to be supported for executorch, could the other two reject or warn when the module has aliased engines? A user who saves as aot_inductor today gets no signal either way, which is the part that seems worth closing regardless of which behavior you pick.

The alias I/O should be supported in both executorch and standard Torch-TensorRT runtime and AOT inductor

You might want to add _declare_aliased_kv_mutations_on_ep to exp program. AoT inductor is more complicated and we will fix that later

Will rebase and enable exported_program on #4459 as well.

Brief warning here for better debugging

This one in particular is an ordinary case. It fires for every engine output
that isn't aliased, which is most of them, so a warning would be hundreds of
lines on a larger model. The legacy path leaves the equivalent skip silent for
the same reason. If you do want logging here, logger.debug rather than a
warning.

here

here

here

The previous comment no longer holds, and debugging can be painful if anything comes wrong in the future. Change the comment to something like

def _reorder_input_names_for_executorch(
    edge_program: ExportedProgram, engine_node: Any, input_names: List[str]
) -> List[str]:
    """Reorder TRT binding names into executorch_call_delegate argument order.

    The runtime binds positionally (``execute()`` arg ``i`` -> input_binding_names
    ``[i]``), but ExecuTorch fusion may permute the delegate placeholders relative
    to the TRT-submodule order that produced ``input_binding_names``. The names
    can't be matched (TRT names are semantic, lowered placeholders are generic
    ``arg_N``), so recover the permutation by node identity: the engine node's
    first arg lists its input nodes in binding order, so sort the names by each
    node's slot among the graph placeholders (its runtime delegate-arg position).

    Only inputs need this. Outputs are also bound positionally, but they are
    ``getitem(engine_node, idx)`` nodes whose index order equals the engine
    output-binding order, and that order survives lowering -- though not because
    the partition is mutation-free. With aliased-I/O (KV-cache) support a TensorRT
    partition *does* produce mutation outputs, and ``arrange_graph_outputs`` does
    move buffer-mutation outputs ahead of user outputs. It stays a no-op here
    because ``_keep_mutated_buffers_above_delegate`` (``partitioner.py``) strips
    the ``delegation_tag`` from mutated buffer placeholders, so they stay out of
    the delegate's state dict and constants; ExecuTorch's ``_get_new_signature``
    then records the mutation as a plain ``USER_OUTPUT`` rather than a
    ``BUFFER_MUTATION`` (it uses the latter only when the delegate itself consumes
    the buffer). The lowered submodule therefore has no mutation specs, so
    ``arrange_graph_outputs`` computes the identity permutation and the getitem
    indices still line up with the engine's output bindings.

    That guarantee is conditional, not structural. If a mutated buffer is ever
    tagged into a delegate, its spec becomes ``BUFFER_MUTATION``, the delegate's
    outputs are permuted, and they would need the same node-identity reordering
    as the inputs below.
    """

The reconstruction drops example_inputs and verifiers, so the pass isn't a drop-in replacement for its input.

ExportedProgram.init also takes example_inputs and verifiers; neither is carried over, so the returned program has them unset. No impact on the current to_edge path, which is presumably why it hasn't surfaced, but it bites in two places:

If the pass moves to the exported_program branch, the saved artifact silently loses its example inputs.
AOTI rejects such a program outright: RuntimeError: exported_program.example_inputs is required to be set in order for AOTInductor compilation. I hit this while testing whether the pass fixes AOTI and had to restore _example_inputs by hand before I could even reach the real failure.
Passing both through costs two keyword arguments.

Fixed in the next push. Both example_inputs and verifiers are threaded
through now.

I should mention that the retrace=False path has a related gap:
create_trt_exp_program constructs its ExportedProgram without
example_inputs either, so AOTI would refuse a program saved that way for the
same reason. That one predates this PR and isn't a dropped value. There's no
source program to carry from, though it does receive arg_inputs it could use.
Let me know if you want that fixed as well.

memory_planned_buffer_device does not exist in the ExecuTorch commit this repo pins (MODULE.bazel:50, 6118688a09, dated 2026-05-18); it landed on release/1.3 after that pin, so //examples/executorch_reference_runner:kv_cache_decode_check does not compile. Bump the pin to a commit that has the API.

Looks like it's already fixed on main: e2f18eb23c (executorch==1.3.1) landed
there in #4398, so this stack picks it up in the rebase. Verified it has the API
but not whether it builds. Will do that.

Whether the mutation gets declared is decided twice, once by use_legacy_exporter inside _exporter.export and once by the outer retrace branch here, so two of the four combinations are wrong: retrace=True, use_legacy_exporter=True declares it twice and raises SpecViolationError: User output getitem_3 is not in the correct order, and retrace=False, use_legacy_exporter=False never declares it and prints no warning. Decide exposure once from the exporter that actually ran, and make _declare_aliased_kv_mutations_on_ep skip a buffer that already has a BUFFER_MUTATION spec.

Turns out declaring on the exported_program branch created a third bad combination on top of your two that you found.

Fixed in the next push, using your suggestion and that covers all three combinations rather than special-casing any of them.

Two cases drop an aliased output here while the runtime still requires one delegate arg per engine output (TensorRTBackend.cpp:426): a second engine aliasing the same buffer is deduped, and a caller-supplied cache that is not a registered buffer gives buf_fqn is None and is skipped, both ending in Error::InvalidArgument at execute. The legacy path skips the same two at lines 765 to 793, so either emit one output arg per engine output, or record in the blob which outputs are not threaded so the runtime can adjust its expected arg count.

I was able to confirm the second case, but I wasn't able to test the first.

I couldn't get a two-engine model that aliases one buffer through the retrace=True path. The export fails earlier, inside torch.export, on the hybrid TRT+eager graph. So I can't say whether that branch is reachable in practice. If you have a model in mind that hits it, that would help.

Regarding the fix, your second option looks workable for the non-buffer case specifically, since there's no ET buffer to write back to. I'd be wary of applying it generally though. An un-threaded aliased output also gets no write-back, and for a real buffer that turns an execute-time error into a silently stale cache. So scoping the blob-side elision to the non-buffer case and keeping one-arg-per-output elsewhere seems safer.

This discards the resize_tensor result and skips the resize entirely on a bad rank, while the non-aliased branch at lines 693 to 706 treats both as fatal. On a dynamic aliased output that outgrows its planned size, nbytes() on the next line is then the stale planned size and drives both the reflect and ExecuTorch's write-back copy_, so handle the return value and the rank check the way the sibling branch does.

dst == bind_ptr cannot happen: the delegate input and output are live at the same time, so the memory planner never places them at the same address, and a mutable buffer is not in the planned arena at all. That makes the zero-copy path described in the comment and in the PR description dead, and every decode step pays a full cache-size device copy here plus ExecuTorch's write-back copy_ on a buffer the engine already updated in place, so it is worth stating whether that cost is inherent to expressing this as BUFFER_MUTATION.

Confirmed dead. I instrumented the branch on a 30-layer export: 120 aliased outputs across prefill and decode, dst == bind_ptr false every time. The comment should go.

One correction: in our .pte the mutable buffers are planned values (mem_id 1, no mutable data segment). The reason dst != bind_ptr is the one you give, plus a second one: the delegate never receives the buffer at all. PropagateDevicePass wraps every delegate input in _h2d_copy, so the engine binds a per-call staging copy, and the reflect plus ExecuTorch's write-back exist to carry the update back to the buffer.

On whether the cost is inherent to BUFFER_MUTATION: it isn't. Pointing the mutation at the buffer placeholder does make ExecuTorch emit no copy_, but that alone breaks persistence, because the staging copy is re-made from the buffer at the top of every call, so the engine's in-place write is overwritten before it is read. Removing the staging too fixes it: the buffer lands in the device arena, the delegate binds it directly, and every full-cache copy goes away. On a single-layer KV model that is 16 -> 8 values and 10 -> 4 instructions, with the persistence check matching eager exactly (0.235384) where the elision-only build reported 0.

That is prototype-grade: a toy model, and the un-staging rule still needs a guard that every consumer of the buffer is device-capable.

The natural home looks like PropagateDevicePass itself, so the copy is never inserted for a delegate-mutated buffer in the first place, though that is more your call than mine. We could also carry it as a Torch-TensorRT pass that undoes the staging for our delegate only, if that sequences better, but it would have to hook in between PropagateDevicePass and memory planning, which is more fragile than doing it properly upstream.

Worth pursuing separately. For this PR I just corrected the comments. I kept the dst != bind_ptr check itself as a self-copy guard rather than deleting it, since it costs a pointer compare and is what stops a self-copy if planning ever changes.

aliased_io changes what a blob means, but the magic stays TR01 and the header is unchanged, so a pre-PR runtime still loads a new blob, ignores the alias map, and binds the KV output to its own allocation instead of the aliased input. I compiled the merge-base parser against head-format metadata and it accepted it, so bump the magic or add a required-feature field the old parser rejects, and let version skew fail closed instead of returning wrong results.

The order these names are serialized in is correct only because _keep_mutated_buffers_above_delegate (partitioner.py:165) keeps the mutated buffer above the delegate, and nothing checks that: inputs get an arity check, outputs get none. An assertion that the lowered delegate's output specs are in engine binding order would turn a future regression from silently swapped outputs into a loud failure.